import typing


# TODO: 不要在 Python 层面实现
# 这里是用 Python 实现了四阶龙格库塔法。
# 实际一般直接使用 C++ 层面的包，否则效率较低。
class LotkaVolterraProblem:
    def __init__(self,
                 a1: float, b1: float,
                 a2: float, b2: float,
                 x1_0: float, x2_0: float,
                 t_step: float = 1E-3):
        if t_step <= 0:
            raise ValueError(
                f"The given argument {t_step=} is less than zero.")
        self._t_step = t_step

        self._a1 = a1
        self._b1 = b1
        self._a2 = a2
        self._b2 = b2

        def f(x1: float, x2: float) -> tuple[float, float]:
            d_x1_d_t = a1 * x1 - b1 * x1 * x2
            d_x2_d_t = a2 * x1 * x2 - b2 * x2
            return d_x1_d_t, d_x2_d_t
        self._f = f

        self.__cached_x1_positive = [x1_0]
        self.__cached_x2_positive = [x2_0]

        self.__cached_x1_negative = [x1_0]
        self.__cached_x2_negative = [x2_0]

    def a1(self) -> float:
        return self._a1

    def b1(self) -> float:
        return self._b1

    def a2(self) -> float:
        return self._a2

    def b2(self) -> float:
        return self._b2

    def x1_0(self) -> float:
        return self.__cached_x1_positive[0]

    def x2_0(self) -> float:
        return self.__cached_x2_positive[0]

    def t_step(self) -> float:
        return self._t_step

    def __solve_to_step(self,
                        step_stop_abs: int,
                        step_stop_sign: int):
        assert step_stop_abs >= 0
        if step_stop_sign >= 0:
            assert step_stop_sign == 1
            x1_caches = self.__cached_x1_positive
            x2_caches = self.__cached_x2_positive
            assert len(x1_caches) == len(x2_caches) <= step_stop_abs
        else:
            assert step_stop_sign == -1
            x1_caches = self.__cached_x1_negative
            x2_caches = self.__cached_x2_negative
            assert len(x1_caches) == len(x2_caches) <= step_stop_abs

        h = self._t_step * step_stop_sign
        h_d_2 = h / 2
        h_d_3 = h / 3
        h_d_6 = h / 6
        for n in range(len(x1_caches) - 1, step_stop_abs):
            p_t = x1_caches[-1]
            p_o = x2_caches[-1]
            k1_t, k1_o = self._f(p_t,
                                 p_o)
            k2_t, k2_o = self._f(p_t + h_d_2 * k1_t,
                                 p_o + h_d_2 * k1_o)
            k3_t, k3_o = self._f(p_t + h_d_2 * k2_t,
                                 p_o + h_d_2 * k2_o)
            k4_t, k4_o = self._f(p_t + h * k3_t,
                                 p_o + h * k3_o)
            x1_caches.append(
                p_t + h_d_6 * k1_t + h_d_3 * k2_t + h_d_3 * k3_t + h_d_6 * k4_t)
            x2_caches.append(
                p_o + h_d_6 * k1_o + h_d_3 * k2_o + h_d_3 * k3_o + h_d_6 * k4_o)

    def x1(self, t: float) -> float:
        step = int(t / self._t_step)
        if step >= 0:
            caches = self.__cached_x1_positive
            sign = 1
        else:
            caches = self.__cached_x1_negative
            step = -step
            sign = -1

        if len(caches) <= step:
            self.__solve_to_step(step, sign)
        return caches[step]

    def x1_batch(self, t_batch: typing.Iterable[float])\
            -> typing.Generator[float, None, None]:
        for t in t_batch:
            yield self.x1(t)

    def x2(self, t: float) -> float:
        step = int(t / self._t_step)
        if step >= 0:
            caches = self.__cached_x2_positive
            sign = 1
        else:
            caches = self.__cached_x2_negative
            step = -step
            sign = -1

        if len(caches) <= step:
            self.__solve_to_step(step, sign)
        return caches[step]

    def x2_batch(self, t_batch: typing.Iterable[float])\
            -> typing.Generator[float, None, None]:
        for t in t_batch:
            yield self.x2(t)
